package com.zhentao.service;

import com.zhentao.config.TableSchemaProperties;
import com.zhentao.util.SqlPromptGenerator;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Service;
import java.util.List;
import java.util.Map;

@Service
public class SqlGenerateService {

    private final OllamaChatClient ollamaChatClient;
    private final JdbcTemplate jdbcTemplate;
    private final SqlPromptGenerator sqlPromptGenerator;
    private final TableSchemaProperties tableSchemaProperties;

    @Autowired
    public SqlGenerateService(OllamaChatClient ollamaChatClient,
                             JdbcTemplate jdbcTemplate,
                             SqlPromptGenerator sqlPromptGenerator,
                             TableSchemaProperties tableSchemaProperties) {
        this.ollamaChatClient = ollamaChatClient;
        this.jdbcTemplate = jdbcTemplate;
        this.sqlPromptGenerator = sqlPromptGenerator;
        this.tableSchemaProperties = tableSchemaProperties;
    }

    /**
     * 将自然语言转换为SQL语句
     */
    public String convertToSql(String naturalLanguage) {
        List<String> tableSchemas = tableSchemaProperties.getTableSchemas();
        if (tableSchemas == null || tableSchemas.isEmpty()) {
            throw new RuntimeException("未配置表结构，请检查application.yml中的sql.table-schemas");
        }
        
        String promptContent = sqlPromptGenerator.generatePromptWithSchemas(tableSchemas, naturalLanguage);
        String generatedSql = ollamaChatClient.call(new Prompt(promptContent))
                .getResult()
                .getOutput()
                .getContent()
                .trim();
        
        // 处理可能的代码块包裹
        if (generatedSql.contains("```sql")) {
            generatedSql = generatedSql.split("```sql")[1].split("```")[0].trim();
        }
        
        // 基本SQL校验
        if (!isValidSql(generatedSql)) {
            throw new RuntimeException("生成的SQL不符合要求: " + generatedSql);
        }
        
        return generatedSql;
    }

    /**
     * 执行SQL语句并返回结果
     */
    public Object executeSql(String sql) {
        try {
            String upperSql = sql.trim().toUpperCase();
            
            if (upperSql.startsWith("SELECT") || upperSql.startsWith("SHOW") || upperSql.startsWith("DESCRIBE")) {
                // 查询语句，返回结果集
                List<Map<String, Object>> result = jdbcTemplate.queryForList(sql);
                return Map.of("type", "query", "data", result, "rows", result.size());
            } else if (upperSql.startsWith("INSERT") || upperSql.startsWith("UPDATE") || upperSql.startsWith("DELETE")) {
                // 增删改语句，返回影响行数
                int rowsAffected = jdbcTemplate.update(sql);
                return Map.of("type", "modify", "rowsAffected", rowsAffected);
            } else {
                throw new RuntimeException("不支持的SQL操作类型: " + sql);
            }
        } catch (Exception e) {
            throw new RuntimeException("SQL执行失败: " + e.getMessage() + ", SQL: " + sql);
        }
    }

    /**
     * 完整流程：自然语言 -> SQL -> 执行结果
     */
    public Map<String, Object> processNaturalLanguage(String naturalLanguage) {
        try {
            String sql = convertToSql(naturalLanguage);
            Object result = executeSql(sql);
            
            return Map.of(
                "success", true,
                "question", naturalLanguage,
                "generatedSql", sql,
                "result", result
            );
        } catch (Exception e) {
            return Map.of(
                "success", false,
                "question", naturalLanguage,
                "error", e.getMessage()
            );
        }
    }

    /**
     * 验证SQL是否为允许的类型
     */
    private boolean isValidSql(String sql) {
        if (sql == null || sql.trim().isEmpty()) {
            return false;
        }
        
        String upperSql = sql.trim().toUpperCase();
        return upperSql.startsWith("SELECT") || 
               upperSql.startsWith("INSERT") || 
               upperSql.startsWith("UPDATE") || 
               upperSql.startsWith("DELETE") ||
               upperSql.startsWith("SHOW") ||
               upperSql.startsWith("DESCRIBE");
    }
}
